# -*- coding:utf-8 -*-
import time
import struct
from socket import *
import os
import sys
import subprocess

import db_server

MYTTL = 255
DEFAULT_MULTICAST_GROUP = '224.1.1.1'
DEFAULT_MULTICAST_PORT = 1234
DEFAULT_SOURCE_ADDR = '127.0.0.1'
DEFAULT_SOURCE_DEV = 'eth0'

class Sender():
    def __init__(self):
        self.multicast_group_addr = DEFAULT_MULTICAST_GROUP
        self.multicast_group_port = DEFAULT_MULTICAST_PORT
        self.source_addr = DEFAULT_SOURCE_ADDR
        self.source_port = 0
        self.source_dev = DEFAULT_SOURCE_DEV
        self.source_data = 1
        self.enable = False

class ExecErr(Exception):
    def __init__(self, returncode, cmd, stdout, stderr):
        self.returncode = returncode
        self.cmd = cmd
        self.stdout = stdout
        self.stderr = stderr

def whiptail(args):
    cmd = ['whiptail', '--output-fd', '2']
    cmd.extend(args)
    p = subprocess.Popen(cmd,
                         shell=False,
                         stdin=sys.stdin,
                         stderr=subprocess.PIPE)
    (stdout, stderr) = p.communicate()
    if p.returncode != 0:
        sys.exit()
    return stderr


# 组播地址范围224.0.0.0 到 239.255.255.255 
def get_multicast_group_addr():
    args = ['--title', '请输入组播组地址(默认地址:'+DEFAULT_MULTICAST_GROUP+')', '--inputbox','', '10', '60']
    multicast_group_addr = DEFAULT_MULTICAST_GROUP
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        multicast_group_addr = str(whiptail_args, encoding='utf-8')
    return multicast_group_addr

# 端口范围 1 - 65535
def get_multicast_group_port():
    args = ['--title', '请输入组播组端口(默认端口:'+str(DEFAULT_MULTICAST_PORT)+')', '--inputbox','', '10', '60']
    multicast_group_port = DEFAULT_MULTICAST_PORT
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        port = int(whiptail_args)
        if port < 1 or port >65535:
            raise ValueError("The multicast group port must be in the range of 1~65535")
        multicast_group_port = port
    return multicast_group_port
# A、B、C类地址
def get_source_addr():
    args = ['--title', '请输入源地址(默认地址:'+DEFAULT_SOURCE_ADDR+')', '--inputbox','', '10', '60']
    source_addr = DEFAULT_SOURCE_ADDR
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        source_addr = str(whiptail_args, encoding='utf-8')
    return source_addr

# 端口范围 1 - 65535
def get_source_port():
    args = ['--title', '请输入源端口(默认端口:随机)', '--inputbox','', '10', '60']
    source_port = 0
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        port = int(whiptail_args)
        if port < 1 or port >65535:
            raise ValueError("The source port must be in the range of 1~65535")
        source_port = port
    return source_port

# 网卡任意填
def get_source_dev():
    args = ['--title', '请输入源网卡(默认网卡:'+DEFAULT_SOURCE_DEV+')', '--inputbox','', '10', '60']
    source_dev = DEFAULT_SOURCE_DEV
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        source_dev = str(whiptail_args, encoding='utf-8')
    return source_dev

# 未定
def get_source_data():
    args = ['--title', '输入源数据包大小(默认包大小:20)', '--inputbox','', '10', '60']
    source_data = 20
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        size = int(whiptail_args)
        if size < 20 or size > 65507:
            raise ValueError("The source size must be in the range of 20~65507")
        source_data = size
    return source_data

def get_enable():
    args = ['--title', '是否开启性能测试', '--menu', '',
        '10', '100', '4',
        'yes', '开启性能测试',
        'no', '不开启性能测试']
    return str(whiptail(args), encoding='utf-8')

def get_sender():
    sender = Sender()
    sender.multicast_group_addr = get_multicast_group_addr()
    sender.multicast_group_port = get_multicast_group_port()
    sender.source_addr = get_source_addr()
    sender.source_port = get_source_port()
    sender.source_dev = get_source_dev()
    sender.source_data = get_source_data()
    if get_enable() == "yes":
        sender.enable = True
    else:
        sender.enable = False
    return sender

def sender():
    sender = get_sender()
    print_info(sender)
    s = socket(AF_INET, SOCK_DGRAM,IPPROTO_UDP)
    s.bind((sender.source_addr, sender.source_port))
    ttl_bin = struct.pack('@i', MYTTL)
    s.setsockopt(IPPROTO_IP, IP_MULTICAST_TTL, ttl_bin)
    s.setsockopt(IPPROTO_IP,
        IP_ADD_MEMBERSHIP,
        inet_aton(sender.multicast_group_addr) + inet_aton(sender.source_addr))
    single_data = '.'
    data = ''
    for num in range(0, sender.source_data):
        data = data +single_data
    num = 0
    start = int(time.time())
    traffic = 0
    oldtraffic = 0
    db_server.insert_udp_multicast_send(sender.multicast_group_addr, sender.multicast_group_port, sender.source_addr, sender.source_port, sender.source_dev, sender.source_data, sender.enable)
    db_server.insert_udp_multicast_receiver(sender.multicast_group_addr, DEFAULT_SOURCE_ADDR, sender.source_port, "unknown")
    udp_multicast_send_id =  db_server.select_udp_multicast_send_id()
    udp_multicast_receiver_id = db_server.select_udp_multicast_receiver_id()
    db_server.insert_udp_multicast_traffic(2, 0, 0, int(time.time()), udp_multicast_receiver_id)
    db_server.insert_udp_multicast_traffic_speed(2, 0, int(time.time()), udp_multicast_receiver_id)
    while True:
        try:
            num = num + 1
            numstr = str(num)
            if num == 1 :
                s.sendto( bytes("start" + data[len("start"):len(data)], encoding='utf8') , (bytes(sender.multicast_group_addr, encoding='utf-8'), sender.multicast_group_port))
            else:
                s.sendto( bytes(numstr + data[len(numstr):len(data)], encoding='utf8') , (bytes(sender.multicast_group_addr, encoding='utf-8'), sender.multicast_group_port))
            now = int(time.time())
            if now - start >=1:
                start = now
                print_info(sender)
                traffic = num*(sender.source_data+28)
                traffic_speed = traffic - oldtraffic
                print_real_info(num, traffic, traffic_speed)
                db_server.insert_udp_multicast_traffic(1, traffic , num , int(time.time()) , udp_multicast_send_id)
                db_server.insert_udp_multicast_traffic_speed(1, traffic_speed , int(time.time()) , udp_multicast_send_id)
            if not sender.enable :
               time.sleep(1)
            oldtraffic = traffic
        except :
            time.sleep(0.5)
            s.sendto(b'close', (bytes(sender.multicast_group_addr, encoding='utf-8'), sender.multicast_group_port))
            print_info(sender)
            traffic = num*(sender.source_data +28 ) + 33
            traffic_speed = 0
            print_real_info(num + 1, traffic, traffic_speed)
            db_server.insert_udp_multicast_traffic(1, traffic , num+1 , int(time.time()) , udp_multicast_send_id)
            db_server.insert_udp_multicast_traffic_speed(1, traffic_speed , int(time.time()) , udp_multicast_send_id)
            if db_server.select_max_receiver_dev() == "unknown":
                db_server.update_receiver_dev()
            s.close()
            print('close sender')
            break
def print_info(sender):
    os.system('clear')
    print("============================")
    print("  组播地址：", sender.multicast_group_addr)
    print("  组播端口：", sender.multicast_group_port)
    print("  源地址：", sender.source_addr)
    if sender.source_port == 0:
        port = '随机'
    else:
        port = sender.source_port
    print("  源端口：", port)
    print("  源网卡：", sender.source_dev)
    print("  源数据包大小：", sender.source_data)
    print("  是否开启性能测试：", sender.enable)
    print("============================")

def print_real_info(num, traffic, traffic_speed):
    print("  发送包的数量：",num)
    print("  发送的流量：",traffic)
    print("  速率：",traffic_speed)

if __name__ == "__main__":
    sender()
